"""
用50-1000-1000-1的全连接网络做脉冲序列学习
"""
import argparse
import torch
from torch import nn
from spikingtorch.seq_learning_utils.seq_learning import *
from spikingtorch.seq_learning_utils.PoissonEncoding import PoissonEncoding
from spikingtorch.seq_learning_utils.CauchySchwarzDistance import CauchySchwarzDistance
from spikingtorch import layer, neuron, base
from tqdm import tqdm

from task.spikingtorch.spikingtorch import dtpdp

class Snn(base.BaseNet):
    def __init__(self, neuronal_decay: float):
        super().__init__()
        self.layers = nn.Sequential(
            layer.Linear(50, 1000, False),
            neuron.LIFNode(tau=neuronal_decay),
            layer.Linear(1000, 100, False),
            neuron.LIFNode(tau=neuronal_decay),
            layer.Linear(100, 1, False),
            neuron.LIFNode(tau=neuronal_decay),
        )

    def forward(self, x: torch.Tensor):
        return self.layers(x)


def spike_sequence_learning(model, learner, args):
    input_spike, target_spike = gen_seq_learning_task(50, rate_input=100, rate_output=100, tsteps=args.T)
    # input_spike = input_spike.unsqueeze(1).to(args.device)
    # target_spike = target_spike.unsqueeze(1).to(args.device)
    input_spike = input_spike.to(args.device)
    target_spike = target_spike.to(args.device)
    print(input_spike.shape)
    print(spike_seq_to_time(input_spike.T[0]))
    print(spike_seq_to_time(input_spike.T[1]))
    print(target_spike.shape)
    print(spike_seq_to_time(target_spike.T[0]))
    max_acc = 0.
    with torch.no_grad():
        for i in range(args.start_epoch, args.epochs):
            out_spike = []
            for t in range(args.T):
                # print(input_spike[t].unsqueeze(0).shape)
                spikes_out = model(input_spike[t].unsqueeze(0)).type(torch.int)
                # print(target_spike[t].shape)
                # exit(0)
                out_spike.append(spikes_out.squeeze())
                learner.step(spikes_out, target_spike[t].unsqueeze(0))
            out_spike = torch.stack(out_spike).unsqueeze(0)
            spike_seq_error = CauchySchwarzDistance(out_spike, target_spike.T)
            acc = spike_seq_error.spike_sequence_error().item()
            if acc > max_acc:
                max_acc = acc
                if args.save_model: torch.save(model.state_dict(), args.save_path+'/model_weights.pth')
            print(out_spike.shape)
            print(f'epoch = {i}, acc = {acc}, max_acc = {max_acc}, output=', spike_seq_to_time(out_spike[0]))
            if args.save_log:
                with open(args.save_path + '/' + args.log_name, 'a') as log_file:
                    log_file.write(f'epoch = {i}, acc = {acc}, max_acc = {max_acc}\n')
            model.reset()
            learner.reset()

def main():
    args = get_hyper_params()
    model = Snn(args.neuronal_decay)
    for name, param in model.named_parameters():
        if 'weight' in name:
            param.data = nn.Parameter(torch.randn_like(param.data))
    model.to(args.device)
    if args.resume: model.load_state_dict(torch.load(args.save_path+'/model_weights.pth'))
    learner = dtpdp.OnlineSupervisedDtrpSTDPLearner(model, lr=args.lr, decay=args.neuronal_decay, Bthres=args.Bthres)
    spike_sequence_learning(model, learner, args)

def get_hyper_params():
    parser = argparse.ArgumentParser()
    parser.add_argument("--T", default=1000, type=int, help="总时长，单位(0.1s)")
    parser.add_argument("--lr", default=1e-3, type=float, help="Learning rate")
    parser.add_argument("--epochs", default=1000, type=int, help="Number of epochs")
    parser.add_argument("--lr-decay", default=4e-2, type=float, help="Exponential decay for learning rate")
    parser.add_argument("--neuronal-thres", default=1.0, type=float, help="Spike threshold for neuronal potential")
    parser.add_argument("--neuronal-decay", default=4.0, type=float, help="Decay rate for membrane potential and spike traces")
    parser.add_argument("--charge-margin", default=2, type=int, help="Tsteps for charging")
    parser.add_argument("--device", default='cuda:0' if torch.cuda.is_available() else 'cpu', type=str)
    parser.add_argument('--Bthres', default=0, type=float, help='随机矩阵B的阈值，低于它则置0')
    parser.add_argument('--save-model', action="store_true", help='保存模型')
    parser.add_argument('--save-path', type=str, help='模型保存位置')
    parser.add_argument('--save-log', action="store_true", help="保存日志")
    parser.add_argument('--log-name', default="snn.log", type=str, help='日志名')
    parser.add_argument('--resume', action='store_true', help='恢复模型')
    parser.add_argument('--start-epoch', default=0, type=int, help='起始epoch')
    args = parser.parse_args()
    return args

if __name__ == '__main__': main()